# Copyright (c) Meta Platforms, Inc. and affiliates.

from .. import CustomSelector, StreamType

class MLModel: ...

class GBTModel(MLModel):
    def __init__(self, model: str) -> None: ...

class FeatureGenerator:
    def __init__(self, feature_names: set[str]) -> None: ...
    def getFeatures(self, data: np.array) -> FeatureMap: ...

class IntFeatureGenerator(FeatureGenerator): ...
class DeltaIntFeatureGenerator(FeatureGenerator): ...
class TokenizeIntFeatureGenerator(FeatureGenerator): ...

class MLSelector(CustomSelector):
    def __init__(
        self,
        input_type: StreamType,
        model: MLModel,
        feature_generator: FeatureGenerator,
        labels: list[str],
    ) -> None: ...

class MemMLTrainingSelector(CustomSelector):
    def __init__(
        self,
        input_type: StreamType,
        labels: list[str],
        collect_inputs: bool = True,
        feature_generator: FeatureGenerator | None = None,
    ) -> None: ...
    def getCollectedJson(self) -> str: ...
    def getCollected(self) -> list[MLTrainingSample]: ...
    def flushCollected(self) -> list[MLTrainingSample]: ...
    def clearCollected(self) -> None: ...

class MLTrainingSample:
    data: np.array | None
    targets: list[int]
    features: list[float]

def samples_from_json(json: str) -> list[MLTrainingSample]: ...
def samples_to_json(samples: list[MLTrainingSample]) -> str: ...
